/**********************************************************************

  Audacity: A Digital Audio Editor

  InterpolateAudio.cpp

  Dominic Mazzoni

**********************************************************************/

#include "InterpolateAudio.h"

#include <math.h>
#include <stdlib.h>

#include <wx/defs.h>

#include "Matrix.h"

static inline int imin(int x, int y)
{
    return x < y ? x : y;
}

static inline int imax(int x, int y)
{
    return x > y ? x : y;
}

// This function is a really dumb, simple way to interpolate audio,
// if the more general InterpolateAudio function below doesn't have
// enough data to work with.  If the bad samples are in the middle,
// it's literally linear.  If it's on either edge, we add some decay
// back to zero.
static void LinearInterpolateAudio(float* buffer, int len,
                                   int firstBad, int numBad)
{
    int i;

    float decay = 0.9f;

    if (firstBad == 0) {
        float delta = buffer[numBad] - buffer[numBad + 1];
        float value = buffer[numBad];
        i = numBad - 1;
        while (i >= 0) {
            value += delta;
            buffer[i] = value;
            value *= decay;
            delta *= decay;
            i--;
        }
    } else if (firstBad + numBad == len) {
        float delta = buffer[firstBad - 1] - buffer[firstBad - 2];
        float value = buffer[firstBad - 1];
        i = firstBad;
        while (i < firstBad + numBad) {
            value += delta;
            buffer[i] = value;
            value *= decay;
            delta *= decay;
            i++;
        }
    } else {
        float v1 = buffer[firstBad - 1];
        float v2 = buffer[firstBad + numBad];
        float value = v1;
        float delta = (v2 - v1) / (numBad + 1);
        i = firstBad;
        while (i < firstBad + numBad) {
            value += delta;
            buffer[i] = value;
            i++;
        }
    }
}

// Here's the main interpolate function, using
// Least Squares AutoRegression (LSAR):
void InterpolateAudio(float* buffer, const size_t len,
                      size_t firstBad, size_t numBad)
{
    const auto N = len;

    wxASSERT(len > 0
             && firstBad >= 0
             && numBad < len
             && firstBad + numBad <= len);

    if (numBad >= len) {
        return; //should never have been called!
    }
    if (firstBad == 0) {
        // The algorithm below has a weird asymmetry in that it
        // performs poorly when interpolating to the left.  If
        // we're asked to interpolate the left side of a buffer,
        // we just reverse the problem and try it that way.
        Floats buffer2{ len };
        for (size_t i=0; i < len; i++) {
            buffer2[len - 1 - i] = buffer[i];
        }
        InterpolateAudio(buffer2.get(), len, len - numBad, numBad);
        for (size_t i=0; i < len; i++) {
            buffer[len - 1 - i] = buffer2[i];
        }
        return;
    }

    Vector s(len, buffer);

    // Choose P, the order of the autoregression equation
    const int IP
        =imin(imin(numBad * 3, 50), imax(firstBad - 1, len - (firstBad + numBad) - 1));

    if (IP < 3 || IP >= (int)N) {
        LinearInterpolateAudio(buffer, len, firstBad, numBad);
        return;
    }

    size_t P(IP);

    // Add a tiny amount of random noise to the input signal -
    // this sounds like a bad idea, but the amount we're adding
    // is only about 1 bit in 16-bit audio, and it's an extremely
    // effective way to avoid nearly-singular matrices.  If users
    // run it more than once they get slightly different results;
    // this is sometimes even advantageous.
    for (size_t i=0; i < N; i++) {
        s[i] += (rand() - (RAND_MAX / 2)) / (RAND_MAX * 10000.0);
    }

    // Solve for the best autoregression coefficients
    // using a least-squares fit to all of the non-bad
    // data we have in the buffer
    Matrix X(P, P);
    Vector b(P);

    for (size_t i = 0; i + P < len; i++) {
        if (i + P < firstBad || i >= (firstBad + numBad)) {
            for (size_t row=0; row < P; row++) {
                for (size_t col=0; col < P; col++) {
                    X[row][col] += (s[i + row] * s[i + col]);
                }
                b[row] += s[i + P] * s[i + row];
            }
        }
    }

    Matrix Xinv(P, P);
    if (!InvertMatrix(X, Xinv)) {
        // The matrix is singular!  Fall back on linear...
        // In practice I have never seen this happen if
        // we add the tiny bit of random noise.
        LinearInterpolateAudio(buffer, len, firstBad, numBad);
        return;
    }

    // This vector now contains the autoregression coefficients
    const Vector& a = Xinv * b;

    // Create a matrix (a "Toeplitz" matrix, as it turns out)
    // which encodes the autoregressive relationship between
    // elements of the sequence.
    Matrix A(N - P, N);
    for (size_t row=0; row < N - P; row++) {
        for (size_t col=0; col < P; col++) {
            A[row][row + col] = -a[col];
        }
        A[row][row + P] = 1;
    }

    // Split both the Toeplitz matrix and the signal into
    // two pieces.  Note that this code could be made to
    // work even in the case where the "bad" samples are
    // not contiguous, but currently it assumes they are.
    //   "u" is for unknown (bad)
    //   "k" is for known (good)
    Matrix Au = MatrixSubset(A, 0, N - P, firstBad, numBad);
    Matrix A_left = MatrixSubset(A, 0, N - P, 0, firstBad);
    Matrix A_right = MatrixSubset(A, 0, N - P,
                                  firstBad + numBad, N - (firstBad + numBad));
    Matrix Ak = MatrixConcatenateCols(A_left, A_right);

    const Vector& s_left = VectorSubset(s, 0, firstBad);
    const Vector& s_right = VectorSubset(s, firstBad + numBad,
                                         N - (firstBad + numBad));
    const Vector& sk = VectorConcatenate(s_left, s_right);

    // Do some linear algebra to find the best possible
    // values that fill in the "bad" area
    Matrix AuT = TransposeMatrix(Au);
    Matrix X1 = MatrixMultiply(AuT, Au);
    Matrix X2(X1.Rows(), X1.Cols());
    if (!InvertMatrix(X1, X2)) {
        // The matrix is singular!  Fall back on linear...
        LinearInterpolateAudio(buffer, len, firstBad, numBad);
        return;
    }
    Matrix X2b = X2 * -1.0;
    Matrix X3 = MatrixMultiply(X2b, AuT);
    Matrix X4 = MatrixMultiply(X3, Ak);
    // This vector contains our best guess as to the
    // unknown values
    const Vector& su = X4 * sk;

    // Put the results into the return buffer
    for (size_t i=0; i < numBad; i++) {
        buffer[firstBad + i] = (float)su[i];
    }
}
